# -*- coding:utf-8 -*-

# @Time    : 2023/4/5 10:56
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : handle_extract.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# 读取xlsx文件,获取标准问及扩展问
import pandas as pd
from util_tool import utils
import re
import copy
from openpyxl import load_workbook


# 根据扩展问知识库，转化成大模型spt的训练语料
def convert_knowledge_to_train_data(path):
    # 打开文件
    datas = pd.read_excel(path)
    # 遍历datas，将每条记录转换成字典，扩展问对于input，标准问对于output
    train_data_list = []
    for index, row in datas.iterrows():
        train_data = {}
        # 获取标准问
        standard_question = row["标准问"]
        # 获取扩展问
        extend_question = row["扩展问"]
        # 判断扩展问为空或nan,或者扩展问长度小于5，跳过
        if (pd.isna(extend_question)) or  (len(extend_question) < 5):
            continue
        train_data["instruction"] = "下面是段销售员和用户的对话，场景是电话销售百万医疗产品，需要对最后一次用户表述进行信息提取。提取的用户意图或用户信息为一个短语"
        train_data["input"] = extend_question
        train_data["output"] =  standard_question
        # 将字典添加到列表中
        train_data_list.append(train_data)
    utils.jdump(train_data_list, "extract_data/341_heiniu_extract.json")
    return train_data_list

# 清理对话内容，只保留文本和标点符号，去除其他符号
def clean_dialogue_content(content):
    # 去掉方括《》及里面的内容
    content = re.sub(r"《.*?》", "", content)
    content = re.sub(r"@#.*?\|\|", "", content)
    content = re.sub(r"@@transfer@@.*?@@notbreak@@}", "", content)
    content = content.replace("#@", "")
    content = content.replace("[第一次静音]", "")
    content = content.replace("[第二次静音]", "")
    content = content.replace("[第三次静音]", "")
    content = content.replace("@no_handle_quiet@", "")
    content = content.replace("@@notbreak@@", "")
    return content

# 清理对话记录，只保留文本和标点符号，去除其他符号
def clean_dialogue(path, save_path):

    wb = load_workbook(path)
    ws = wb[wb.sheetnames[0]]

    datas = pd.read_excel(path)
    dialogue_list = []
    for index, row in datas.iterrows():
        record = {}
        session_id = row["sessionId"]
        role = row["角色"]
        content = row["过滤文本"]
        if role == "销售员":
            content = clean_dialogue_content(content)

        record["session_id"] = session_id
        record["role"] = role
        record["content"] = content
        dialogue_list.append(record)
    # dialogue_list 转成pandas
    df = pd.DataFrame(dialogue_list)
    # 保存到csv文件
    df.to_csv(save_path, index=False, encoding="utf-8-sig")

def _format_dialogue(self, instruction: str, response: str = ''):
    if not response:
        return f'\n\n### Input:\n{instruction}'
    return f'\n\n### Input:\n{instruction}\n\n### Response:\n{response}'

# 讲对话记录转成输入字符串
def convert_dialogue_to_input_string(record_list, max_len=1000):
    input_string = ""
    # 倒序遍历record_list
    for record in record_list[::-1]:
        if record["role"] == "用户":
            input_string = "用户:" + record["content"] + "\n" + input_string
        elif record["role"] == "机器人":
            input_string = "销售员:" + (re.sub(r'\[[^\]]*\]', "", record["content"])).strip() + "\n" + input_string
        if len(input_string) > max_len:
            break
    return input_string


# 将对话记录转成大模型spt的训练语料
def convert_dialogue_to_train_data(path, save_path):
    datas = pd.read_csv(path)
    dialogue_list = []
    dialogue = []
    session_id = ""
    for index, row in datas.iterrows():
        if (row["session_id"] != session_id):
            if len(dialogue) > 0:
                dialogue_list.append(dialogue)
                dialogue = []
            session_id = row["session_id"]
            dialogue.append(row)
        else:
            dialogue.append(row)
    # 将dialogue_list转成train_data_list
    train_data_list = []
    for dialogue in dialogue_list:
        train_data = {}
        train_data["instruction"] = "下面是段销售员和用户的对话，场景是电话销售百万医疗产品，需要对最后一次用户表述进行信息提取。提取的用户意图或用户信息为一个短语"
        input_list = []
        output = ""
        for record in dialogue:
            if record["role"] == "用户":
                input_list.append(record)
            elif record["role"] == "机器人":
                content = record["content"]
                content = re.sub(r'\[发短信:.*?\]', "", content)
                # 若[.*:.*]在content中
                if re.search(r'\[.*?:.*?\]', content) and len(input_list) > 0:
                    # 获取[.*:.*]中的内容
                    match_list = re.findall(r'\[.*?.*?\]', content)
                    # 通过,拼接match_list中所有包含:的内容
                    match_real_list = []
                    for match in match_list:
                        match = match.replace('[', "")
                        match = match.replace("]", "")
                        if ("坐席询问用户" in match) and ("回应" in match):
                            match = match.replace("用户表示", "用户")
                        if ("用户表示不需要" in match):
                            match = "用户主动表示:不需要-无原因"
                        if ("用户询问什么平台" in match):
                            match = "用户主动询问:什么平台"
                        if ("用户询问什么产品" in match):
                            match = "用户主动询问:什么产品"
                        if (":" in match) and ("用户命中FAQ" not in match):
                            match = match.replace("坐席询问用户:FAQ", "坐席询问用户:")

                            match_real_list.append(match)
                    if match_real_list:
                        output = ",".join(match_real_list)
                        train_data["input"] = convert_dialogue_to_input_string(input_list)
                        train_data["output"] = output
                        print(train_data)
                        train_data_list.append(copy.deepcopy(train_data))

                input_list.append(record)
    utils.jdump(train_data_list, save_path)




if __name__ == '__main__':
    # train_data_list = convert_knowledge_to_train_data("extract_data/341-黑牛.xlsx")
    origin_dialogue_path = "dialogue_data/_2199_2023-03-30 00_00_00_2023-03-30 23_10_00_100000000.xlsx"
    clean_save_path = "dialogue_data/2199_cleaned_dialogue.csv"
    dalogue_save_path = "dialogue_data/2199_dialogue_extract.json"
    # origin_dialogue_path = "extract_data/_2023-04-01 00_09_00-2023-04-05 23_59_00-600000.xlsx"
    # clean_save_path = "extract_data/heiniu_cleaned_dialogue_04-05.csv"
    # dalogue_save_path = "extract_data/heiniu_dialogue_extract_0405.json"
    clean_dialogue(origin_dialogue_path, clean_save_path)
    convert_dialogue_to_train_data(clean_save_path, dalogue_save_path)